from datetime import datetime
import os


class ScreenshotUtils:
    """截图工具类（优化版）"""

    def __init__(self, driver, screenshot_dir="test_screenshots"):
        self.driver = driver
        self.screenshot_dir = screenshot_dir
        self._ensure_directory_exists()
        print(f"截图将保存到: {os.path.abspath(self.screenshot_dir)}")

    def _ensure_directory_exists(self):
        """确保截图目录存在（优化性能）"""
        if not os.path.exists(self.screenshot_dir):
            os.makedirs(self.screenshot_dir)

    def take_screenshot(self, screenshot_name, include_timestamp=True):
        """
        截取屏幕截图（返回绝对路径）
        """
        if include_timestamp:
            timestamp = datetime.now().strftime("%H%M%S_%f")[:-3]  # 简化时间戳
            filename = f"{screenshot_name}_{timestamp}.png"
        else:
            filename = f"{screenshot_name}.png"

        filepath = os.path.join(self.screenshot_dir, filename)
        self.driver.save_screenshot(filepath)

        # 打印绝对路径方便查找
        abs_path = os.path.abspath(filepath)
        print(f"截图已保存: {abs_path}")
        return abs_path  # 返回绝对路径

    def take_success_screenshot(self, test_name):
        """成功时截图"""
        filepath = self.take_screenshot(f"SUCCESS_{test_name}")
        print(f"✓ 测试成功: {filepath}")
        return filepath

    def take_failure_screenshot(self, test_name):
        """失败时截图"""
        filepath = self.take_screenshot(f"FAILURE_{test_name}")
        print(f"✗ 测试失败: {filepath}")
        return filepath

    def take_step_screenshot(self, step_name):
        """步骤截图"""
        return self.take_screenshot(f"STEP_{step_name}")


# 简化使用函数
def capture_screenshot(driver, name, dir_name="screenshots"):
    """快速截图函数"""
    utils = ScreenshotUtils(driver, dir_name)
    return utils.take_screenshot(name)